import cv2
import face_recognition
import pickle
import os
import numpy as np
from facenet_training_face_recognition import FaceRecognitionTrainer

def main():
    """
    主函数：在树莓派上运行实时人脸识别
    """
    print("树莓派人脸识别系统启动中...")
    
    
    # 加载训练好的模型
    try:
        user = FaceRecognitionTrainer()
        user.load_model()
        
        #classifier, known_names = load_model("face_recognition_model.pkl")
        if (user.classifier!=None and user.known_names!=None):print("人脸识别模型加载成功")
        else:print("人脸识别模型加载失败")
    except Exception as e:
        print(f"模型加载失败: {e}")
        return
    
    # 初始化摄像头
    video_capture = cv2.VideoCapture(0)
    
    # 检查摄像头是否成功打开
    if not video_capture.isOpened():
        print("错误: 无法打开摄像头")
        return
    
    # 设置摄像头分辨率 (可根据树莓派性能调整)
    video_capture.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    video_capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    
    print("人脸识别系统已启动")
    print("按 'q' 键退出程序")
    
    try:
        while True:
            # 读取一帧图像
            ret, frame = video_capture.read()
            
            if not ret:
                print("无法获取摄像头画面")
                break
            
            # 识别人脸
            result = user.predict(frame)
            if result is not None:
                face_locations, face_names, confidences = result
                # 检查是否检测到人脸
                if face_locations is None:
                    # 未检测到人脸，显示原始帧
                    cv2.imshow('Raspberry Pi Face Recognition', frame)
                    # 按 'q' 键退出
                    if cv2.waitKey(1) & 0xFF == ord('q'):
                        break
                    continue
            else:
                continue
            
            # 在图像上绘制人脸框和标签
            for (top, right, bottom, left), name, confidence in zip(face_locations, face_names, confidences):
                # 绘制人脸框
                cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)
                
                # 绘制标签背景
                cv2.rectangle(frame, (left, bottom - 35), (right, bottom), (0, 0, 255), cv2.FILLED)
                font = cv2.FONT_HERSHEY_DUPLEX
                
                # 绘制人名标签
                cv2.putText(frame, name, (left + 6, bottom - 6), font, 0.8, (255, 255, 255), 1)
                # 绘制置信度
                cv2.putText(frame, f"Conf: {confidence:.2f}", (left + 6, bottom - 20), font, 0.5, (255, 255, 255), 1)
            
            # 显示结果图像
            cv2.imshow('Raspberry Pi Face Recognition', frame)
            
            # 按 'q' 键退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
                
    except KeyboardInterrupt:
        print("\n用户中断程序")
    finally:
        # 释放资源
        video_capture.release()
        cv2.destroyAllWindows()
        print("程序已退出")

if __name__ == "__main__":
    main()